/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.pig.shock;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
import java.net.ConnectException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.net.Socket;
import java.net.SocketAddress;
import java.net.SocketException;
import java.net.SocketImpl;
import java.net.SocketOptions;
import java.net.SocketImplFactory;
import java.net.UnknownHostException;
import java.net.Proxy.Type;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.Properties;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.jcraft.jsch.ChannelDirectTCPIP;
import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Logger;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.SocketFactory;
import com.jcraft.jsch.UserInfo;

/**
 * This class replaces the standard SocketImplFactory with a factory
 * that will use an SSH proxy. Only connect operations are supported. There
 * are no server operations and nio.channels are not supported.
 * <p>
 * This class uses the following system properties:
 * <ul>
 * <li>user.home - used to calculate the default .ssh directory
 * <li>user.name - used for the default ssh user id
 * <li>ssh.user - will override user.name as the user name to send over ssh
 * <li>ssh.gateway - sets the ssh host that will proxy connections
 * <li>ssh.knownhosts - the known hosts file. (We will not add to it)
 * <li>ssh.identity - the location of the identity file. The default name
 *       is openid in the ssh directory. THIS FILE MUST NOT BE PASSWORD
 *       PROTECTED.
 * </ul>
 * @author breed
 *
 */
public class SSHSocketImplFactory implements SocketImplFactory, Logger {
    
    private static final Log log = LogFactory.getLog(SSHSocketImplFactory.class);

	Session session;

	public static SSHSocketImplFactory getFactory() throws JSchException, IOException {
		return getFactory(System.getProperty("ssh.gateway"));
	}
	static HashMap<String, SSHSocketImplFactory> factories = new HashMap<String, SSHSocketImplFactory>();
	public synchronized static SSHSocketImplFactory getFactory(String host) throws JSchException, IOException {
		SSHSocketImplFactory factory = factories.get(host);
		if (factory == null) {
			factory = new SSHSocketImplFactory(host);
			factories.put(host, factory);
		}
		return factory;
	}
	
	private SSHSocketImplFactory(String host) throws JSchException, IOException {
		JSch jsch = new JSch();
		jsch.setLogger(this);
		String passphrase = "";
		String defaultSSHDir = System.getProperty("user.home") + "/.ssh";
		String identityFile = defaultSSHDir + "/openid";
		String user = System.getProperty("user.name");
		user = System.getProperty("ssh.user", user);
		if (host == null) {
			throw new RuntimeException(
					"ssh.gateway system property must be set");
		}
		String knownHosts = defaultSSHDir + "/known_hosts";
		knownHosts = System.getProperty("ssh.knownhosts", knownHosts);
		jsch.setKnownHosts(knownHosts);
		identityFile = System.getProperty("ssh.identity", identityFile);
		jsch.addIdentity(identityFile, passphrase.getBytes());
		session = jsch.getSession(user, host);
		Properties props = new Properties();
		props.put("compression.s2c", "none");
		props.put("compression.c2s", "none");
		props.put("cipher.s2c", "blowfish-cbc,3des-cbc");
		props.put("cipher.c2s", "blowfish-cbc,3des-cbc");
		if (jsch.getHostKeyRepository().getHostKey(host, null) == null) {
			// We don't have a way to prompt, so if it isn't there we want
			// it automatically added.
			props.put("StrictHostKeyChecking", "no");
		}
		session.setConfig(props);
		session.setDaemonThread(true);

		// We have to make sure that SSH uses it's own socket factory so
		// that we don't get recursion
		SocketFactory sfactory = new SSHSocketFactory();
		session.setSocketFactory(sfactory);
		UserInfo userinfo = null;
		session.setUserInfo(userinfo);
		session.connect();
		if (!session.isConnected()) {
			throw new IOException("Session not connected");
		}
	}

	public SocketImpl createSocketImpl() {
		return new SSHSocketImpl(session);

	}

	public boolean isEnabled(int arg0) {
		// Default to not logging anything
		return false;
	}

	public void log(int arg0, String arg1) {
		log.error(arg0 + ": " + arg1);
	}

	class SSHProcess extends Process {
		ChannelExec channel;

		InputStream is;

		InputStream es;

		OutputStream os;

		SSHProcess(ChannelExec channel) throws IOException {
			this.channel = channel;
			is = channel.getInputStream();
			es = channel.getErrStream();
			os = channel.getOutputStream();
		}

		/* (non-Javadoc)
		 * @see java.lang.Process#destroy()
		 */
		@Override
		public void destroy() {
			channel.disconnect();
		}

		/* (non-Javadoc)
		 * @see java.lang.Process#exitValue()
		 */
		@Override
		public int exitValue() {
			return channel.getExitStatus();
		}

		/* (non-Javadoc)
		 * @see java.lang.Process#getErrorStream()
		 */
		@Override
		public InputStream getErrorStream() {
			return es;
		}

		/* (non-Javadoc)
		 * @see java.lang.Process#getInputStream()
		 */
		@Override
		public InputStream getInputStream() {
			return is;
		}

		/* (non-Javadoc)
		 * @see java.lang.Process#getOutputStream()
		 */
		@Override
		public OutputStream getOutputStream() {
			return os;
		}

		/* (non-Javadoc)
		 * @see java.lang.Process#waitFor()
		 */
		@Override
		public int waitFor() throws InterruptedException {
			while (channel.isConnected()) {
				Thread.sleep(1000);
			}
			return channel.getExitStatus();
		}

	}

	public Process ssh(String cmd) throws JSchException, IOException {
		ChannelExec channel = (ChannelExec) session.openChannel("exec");
		channel.setCommand(cmd);
		channel.setPty(true);
		channel.connect();
		return new SSHProcess(channel);
	}
}

/**
 * This socket factory is only used by SSH. We implement it using nio.channels
 * since those classes do not use the SocketImplFactory.
 * 
 * @author breed
 *
 */
class SSHSocketFactory implements SocketFactory {

    private final static Log log = LogFactory.getLog(SSHSocketFactory.class);
    
	public Socket createSocket(String host, int port) throws IOException,
			UnknownHostException {
		String socksHost = System.getProperty("socksProxyHost");
		Socket s;
		InetSocketAddress addr = new InetSocketAddress(host, port);
		if (socksHost != null) {
			Proxy proxy = new Proxy(Type.SOCKS, new InetSocketAddress(
					socksHost, 1080));
			s = new Socket(proxy);
			s.connect(addr);
		} else {
			log.error(addr);
			SocketChannel sc = SocketChannel.open(addr);
			s = sc.socket();
		}
		s.setTcpNoDelay(true);
		return s;
	}

	public InputStream getInputStream(Socket socket) throws IOException {
		return new ChannelInputStream(socket.getChannel());
	}

	public OutputStream getOutputStream(Socket socket) throws IOException {
		return new ChannelOutputStream(socket.getChannel());
	}
}

class ChannelOutputStream extends OutputStream {
	SocketChannel sc;

	public ChannelOutputStream(SocketChannel sc) {
		this.sc = sc;
	}

	@Override
	public void write(int b) throws IOException {
		byte bs[] = new byte[1];
		bs[0] = (byte) b;
		write(bs);
	}

	@Override
	public void write(byte b[], int off, int len) throws IOException {
		sc.write(ByteBuffer.wrap(b, off, len));
	}

}

class ChannelInputStream extends InputStream {
	SocketChannel sc;

	public ChannelInputStream(SocketChannel sc) {
		this.sc = sc;
	}

	@Override
	public int read() throws IOException {
		byte b[] = new byte[1];
		if (read(b) != 1) {
			return -1;
		}
		return b[0] & 0xff;
	}

	@Override
	public int read(byte b[], int off, int len) throws IOException {
		return sc.read(ByteBuffer.wrap(b, off, len));
	}
}

/**
 * We aren't going to actually create any new connection, we will forward
 * things to SSH.
 */
class SSHSocketImpl extends SocketImpl {
    
    private static final Log log = LogFactory.getLog(SSHSocketImpl.class);

	Session session;

	ChannelDirectTCPIP channel;

	InputStream is;

	OutputStream os;

	SSHSocketImpl(Session session) {
		this.session = session;
	}

	@Override
	protected void accept(SocketImpl s) throws IOException {
		throw new IOException("SSHSocketImpl does not implement accept");
	}

	@Override
	protected int available() throws IOException {
		if (is == null) {
			throw new ConnectException("Not connected");
		}
		return is.available();
	}

	@Override
	protected void bind(InetAddress host, int port) throws IOException {
		if ((host != null && !host.isAnyLocalAddress()) || port != 0) {
			throw new IOException("SSHSocketImpl does not implement bind");
		}
	}

	@Override
	protected void close() throws IOException {
		if (channel != null) {
			//	    channel.disconnect();
			is = null;
			os = null;
		}
	}

	public final static String defaultDomain = ".inktomisearch.com";

	@Override
	protected void connect(String host, int port) throws IOException {
		InetAddress addr = null;
		try {
			addr = InetAddress.getByName(host);
		} catch (UnknownHostException e) {
			host += defaultDomain;
			addr = InetAddress.getByName(host);
		}
		connect(addr, port);
	}

	@Override
	protected void connect(InetAddress address, int port) throws IOException {
		connect(new InetSocketAddress(address, port), 300000);
	}

	@Override
	protected void connect(SocketAddress address, int timeout)
			throws IOException {
		try {
			if (!session.isConnected()) {
				session.connect();
			}
			channel = (ChannelDirectTCPIP) session.openChannel("direct-tcpip");
			//is = channel.getInputStream();
			//os = channel.getOutputStream();
			channel.setHost(((InetSocketAddress) address).getHostName());
			channel.setPort(((InetSocketAddress) address).getPort());
			channel.setOrgPort(22);
			is = new PipedInputStream();
			os = new PipedOutputStream();
			channel
					.setInputStream(new PipedInputStream((PipedOutputStream) os));
			channel
					.setOutputStream(new PipedOutputStream(
							(PipedInputStream) is));
			channel.connect();
			if (!channel.isConnected()) {
				log.error("Not connected");
			}
			if (channel.isEOF()) {
				log.error("EOF");
			}
		} catch (JSchException e) {
			log.error(e);
			IOException newE = new IOException(e.getMessage());
			newE.setStackTrace(e.getStackTrace());
			throw newE;
		}
	}

	@Override
	protected void create(boolean stream) throws IOException {
		if (stream == false) {
			throw new IOException("Cannot handle UDP streams");
		}
	}

	@Override
	protected InputStream getInputStream() throws IOException {
		return is;
	}

	@Override
	protected OutputStream getOutputStream() throws IOException {
		return os;
	}

	@Override
	protected void listen(int backlog) throws IOException {
		throw new IOException("SSHSocketImpl does not implement listen");
	}

	@Override
	protected void sendUrgentData(int data) throws IOException {
		throw new IOException("SSHSocketImpl does not implement sendUrgentData");
	}

	public Object getOption(int optID) throws SocketException {
        if (optID == SocketOptions.SO_SNDBUF)
            return new Integer(1024);
        else
		    throw new SocketException("SSHSocketImpl does not implement getOption for " + optID);
	}

	/**
	 * We silently ignore setOptions because they do happen, but there is
	 * nothing that we can do about it.
	 */
	public void setOption(int optID, Object value) throws SocketException {
	}

	static public void main(String args[]) {
		try {
			System.setProperty("ssh.gateway", "ucdev2");
			final SSHSocketImplFactory fac = SSHSocketImplFactory.getFactory();
			Socket.setSocketImplFactory(fac);
			for (int i = 0; i < 10; i++) {
				new Thread() {
					@Override
					public void run() {
						try {
							log.error("Starting " + this);
							connectTest("www.yahoo.com");
							log.error("Finished " + this);
						} catch (Exception e) {
							log.error(e);
						}
					}
				}.start();
			}
			Thread.sleep(1000000);
			connectTest("www.news.com");
			log.info("******** Starting PART II");
			for (int i = 0; i < 10; i++) {
				new Thread() {
					@Override
					public void run() {
						try {
							log.error("Starting " + this);
							connectTest("www.flickr.com");
							log.error("Finished " + this);
						} catch (Exception e) {
							log.error(e);
						}
					}
				}.start();
			}
		} catch (Exception e) {
			log.error(e);
		}
	}

	private static void connectTest(String host) throws JSchException,
			IOException {
		Socket s = new Socket(host, 80);
		s.getOutputStream().write("GET / HTTP/1.0\r\n\r\n".getBytes());
		byte b[] = new byte[80];
		int rc = s.getInputStream().read(b);
		System.out.write(b, 0, rc);
		s.close();
	}

	private static void lsTest(SSHSocketImplFactory fac) throws JSchException,
			IOException {
		Process p = fac.ssh("ls");
		byte b[] = new byte[1024];
		final InputStream es = p.getErrorStream();
		new Thread() {
			@Override
			public void run() {
				try {
					while (es.available() > 0) {
						es.read();
					}
				} catch (Exception e) {
				}
			}
		}.start();
		p.getOutputStream().close();
		InputStream is = p.getInputStream();
		int rc;
		while ((rc = is.read(b)) > 0) {
			System.out.write(b, 0, rc);
		}
	}
}
